import numpy as np
from scipy import sparse
import os
import socket

gridN = 66

ku_grid = np.linspace(-1, 1, gridN)
# ku_grid = np.array([-1, 0])

ky_grid = np.linspace(-1, 1, gridN)

kv_grid = np.linspace(-1, 1, gridN)
# kv_grid = np.array([-1, 0])

# qy_measure_data = np.load(r"D:\vc\Rb87PythonData\3DSOCbandData\190618\
# V0=2.000,Vy=-4.000,M=1.155,mz=-0.200_qy_measure.npz")
# qy_measure_list = qy_measure_data['qy_measure_list']
# ky_grid = np.concatenate((ky_grid, qy_measure_list))

Fixed_axis = ku_grid[0]

# kz_grid = np.array([-1, 0, 1])
kz_grid = np.linspace(-1, 1, gridN)
kz_grid = np.delete(kz_grid, (0, gridN - 1))
kx_grid_original = np.linspace(-1, 1, gridN)

# mz_list = np.linspace(-1, 1, 1001)
# mz_list = np.array([-200])
mz_list = np.array([0.127, 0.131, 0.145, 0.168, 0.199, 0.237, 0.282, 0.332,
                    0.385, 0.442, 0.5, 0.558, 0.615, 0.67, 0.72, 0.765, 0.803, 0.833, 0.856, 0.872, 0.884]) / 2

# psi_list = np.linspace(-np.pi, np.pi, 101)
psi_list = np.array([np.pi])

date = '191107'
version_count = 1

Host_Name = socket.gethostname()
if Host_Name == 'VictorChengPC' or Host_Name == 'VictorChengXPS':
    Puter_Name_Pre = 'X:\\work'
    save_address = Puter_Name_Pre + '\\Rb87PythonData\\3DSOCbandData' + '\\' + date + '\\'
    save_pic_address = Puter_Name_Pre + '\\Rb87PythonData\\3DSOCbandData\\pics' + '\\' + date + '\\'
elif Host_Name == 'calcenter2':
    Puter_Name_Pre = 'D:\\vc'
    save_address = Puter_Name_Pre + '\\Rb87PythonData\\3DSOCbandData' + '\\' + date + '\\'
    save_pic_address = Puter_Name_Pre + '\\Rb87PythonData\\3DSOCbandData\\pics' + '\\' + date + '\\'
elif Host_Name == 'Rb87BEC3':
    Puter_Name_Pre = '/media/rb87bec3/Files/vc'
    save_address = Puter_Name_Pre + '/Rb87PythonData/3DSOCbandData' + '/' + date + '/'
    save_pic_address = Puter_Name_Pre + '/Rb87PythonData/3DSOCbandData/pics' + '/' + date + '/'

if os.path.exists(save_address) == 0:
    os.system('mkdir ' + save_address)
if os.path.exists(save_pic_address) == 0:
    os.system('mkdir ' + save_pic_address)

num_of_band = 2

# del_psi = 0
V0 = -1.77  # 1.77
Omega0 = 1.02210 * 2 / np.sqrt(2)
Vy = -3.54  # 3.98

# Vy_List = np.linspace(0, -8, 501)
Vy_List = np.array([Vy])

# Vxz = V0*np.cos(del_psi)
# Omega_xy = Omega0*np.sin(del_psi/2+np.pi/4)
# Omega_zy = Omega0*np.cos(del_psi/2+np.pi/4)
# t1 = Vxz
# t2 = Vy
# t3 = Omega_xy
# t4 = Omega_zy

# Cal_TBM_Par_file = \
# Puter_Name_Pre+r"\Rb87PythonData\3DSOCbandData\190328\3x3x3;4x81x81x81;V_xz=2.000E_r,V_y=-2.000E_r,Oxy=0.800E_r,Ozy=-0.800E_r,T=1.000e-07K,n=3.000e+18m-3del_psi=1.000pi,;mz=0.000Er_v4.npz"
# (t0, ty, tSO) = Cal_TBM_Parameters.cal(Cal_TBM_Par_file)
(t0, ty, tSO) = (0, 0, 0)

AtomTemperture = 100e-9
AtomDensityHot = 3e18
hbar = 1.0545718e-34
kB = 1.38064852e-23
WaveLength = 786.99e-9
Mass = 1.443160648e-25
Er = np.power(hbar * 2 * np.pi / WaveLength, 2) / (2 * Mass)

# kx_grid = np.linspace(-1, 1, 81)
# kz_grid = np.array([-1, 0, 1])

Nu = ku_grid.size
Ny = ky_grid.size
Nv = kv_grid.size

# Nz = kz_grid.size

m0 = 3
n0 = 3
l0 = 3
dim = (2 * m0 + 1) * (2 * n0 + 1) * (2 * l0 + 1) * 2

kus = np.arange(-m0, m0 + 1, 1)
kys = np.arange(-n0, n0 + 1, 1)
kvs = np.arange(-l0, l0 + 1, 1)

Ma = 1
# delta_psi = np.pi
# V0 = 4
# Omega_0 = np.sqrt(2)
# Vxz = V0*np.cos(delta_psi)
# Vy = V0
# Omega_xy = Omega_0*np.sin(delta_psi/2+np.pi/4)
# Omega_zy = Omega_0*np.cos(delta_psi/2+np.pi/4)
# Right config: V0 is neg, Omega0 is pos, del_psi=pi


sigma_0 = np.eye(2)
sigma_x = np.array([[0, 1], [1, 0]])
sigma_y = np.array([[0, -1j], [1j, 0]])
sigma_z = np.array([[1, 0], [0, -1]])
sigma_1 = np.array([[1, 0], [0, 0]])
sigma_2 = np.array([[0, 1], [0, 0]])
sigma_3 = np.array([[0, 0], [1, 0]])
sigma_4 = np.array([[0, 0], [0, 1]])

one_m = sparse.dia_matrix((np.diag(np.eye(2 * m0 + 1)), 0), shape=(2 * m0 + 1, 2 * m0 + 1))
one_n = sparse.dia_matrix((np.diag(np.eye(2 * n0 + 1)), 0), shape=(2 * n0 + 1, 2 * n0 + 1))
one_l = sparse.dia_matrix((np.diag(np.eye(2 * l0 + 1)), 0), shape=(2 * l0 + 1, 2 * l0 + 1))

# temp = np.zeros((2*m0, 1))
mp1_m_matrix = sparse.dia_matrix((np.diag(np.eye(2 * m0 + 1)), -1), shape=(2 * m0 + 1, 2 * m0 + 1))
np1_n_matrix = sparse.dia_matrix((np.diag(np.eye(2 * n0 + 1)), -1), shape=(2 * n0 + 1, 2 * n0 + 1))
lp1_l_matrix = sparse.dia_matrix((np.diag(np.eye(2 * l0 + 1)), -1), shape=(2 * l0 + 1, 2 * l0 + 1))
# input()
